import datetime
import json
import os
import ssl
import sys
import time
import urllib
from typing import Dict, Optional

import certifi
import stix2
import yaml
from pycti import (
    Identity,
    Infrastructure,
    OpenCTIConnectorHelper,
    StixCoreRelationship,
    Vulnerability,
    get_config_variable,
)


class Cisa:
    def __init__(self):
        # Instantiate the connector helper from config
        config_file_path = os.path.dirname(os.path.abspath(__file__)) + "/config.yml"
        config = (
            yaml.load(open(config_file_path), Loader=yaml.FullLoader)
            if os.path.isfile(config_file_path)
            else {}
        )
        self.helper = OpenCTIConnectorHelper(config)
        # Extra config
        self.cisa_catalog_url = get_config_variable(
            "CISA_CATALOG_URL", ["cisa", "catalog_url"], config
        )
        self.cisa_create_infrastructures = get_config_variable(
            "CISA_CREATE_INFRASTRUCTURES",
            ["cisa", "create_infrastructures"],
            config,
            True,
        )
        self.cisa_interval = get_config_variable(
            "CISA_INTERVAL", ["cisa", "interval"], config, True
        )
        self.tlp = get_config_variable(
            "CISA_TLP", ["cisa", "tlp"], config, False, "TLP:CLEAR"
        )
        self.update_existing_data = get_config_variable(
            "CONNECTOR_UPDATE_EXISTING_DATA",
            ["connector", "update_existing_data"],
            config,
        )
        self.confidence_level = get_config_variable(
            "CONNECTOR_CONFIDENCE_LEVEL",
            ["connector", "confidence_level"],
            config,
        )
        self.created_by_stix = None
        self.tlp_marking = None
        self.org = "Cybersecurity and Infrastructure Security Agency"
        self.opencti_url = get_config_variable(
            "OPENCTI_URL", ["opencti", "url"], config
        )
        self.opencti_token = get_config_variable(
            "OPENCTI_TOKEN", ["opencti", "token"], config
        )

    def get_interval(self):
        return int(self.cisa_interval) * 60 * 60 * 24

    def retrieve_data(self, url: str) -> Optional[str]:
        """
        Retrieve data from the given url.

        Parameters
        ----------
        url : str
            Url to retrieve.

        Returns
        -------
        str
            A string with the content or None in case of failure.
        """
        try:
            return (
                urllib.request.urlopen(
                    url,
                    context=ssl.create_default_context(cafile=certifi.where()),
                )
                .read()
                .decode("utf-8")
            )
        except (
            urllib.error.URLError,
            urllib.error.HTTPError,
            urllib.error.ContentTooShortError,
        ) as urllib_error:
            self.helper.log_error(f"Error retrieving url {url}: {urllib_error}")
        return None

    # Get Identity info
    def set_created_by_stix(self, org: str) -> Dict:
        if org is None:
            org = self.org
        org_stix = stix2.Identity(
            id=Identity.generate_id(org, "organization"),
            identity_class="organization",
            name=f"{org}",
            description="The Cybersecurity and Infrastructure Security Agency is a United States federal agency, an operational component under Department of Homeland Security oversight. Its activities are a continuation of the National Protection and Programs Directorate.",
            allow_custom=True,
            custom_properties={"x_opencti_organization_type": "vendor"},
        )
        self.created_by_stix = org_stix

    def create_relationship_obj(
        self, source_ref: str, target_ref: str
    ) -> stix2.Relationship:
        relationship_stix = stix2.Relationship(
            id=StixCoreRelationship.generate_id("has", source_ref, target_ref),
            relationship_type="has",
            source_ref=source_ref,
            target_ref=target_ref,
        )
        return relationship_stix

    def set_tlp_marking(self, tlp_mark):
        if tlp_mark is None:
            tlp_mark = self.tlp
        self.helper.log_info("Retrieving TLP Data from CTI Service")
        marking = self.helper.api.marking_definition.read(
            filters=[{"key": "definition", "values": [f"{tlp_mark}"]}]
        )
        self.tlp_marking = marking["standard_id"]
        self.helper.log_info(f"Marking Definition: {self.tlp_marking}")

    def build_bundle(self, data):
        self.helper.log_info("Building CISA Bundle")
        vuln = data
        stix_objects = [self.created_by_stix]
        vuln_cve = vuln["cveID"]
        vendor_name = vuln["vendorProject"]
        product = vuln["product"]
        description = vuln["shortDescription"]
        vuln_date = vuln["dateAdded"]
        created = f"{vuln_date}T00:00:00.000Z"
        created_by_id = self.created_by_stix["id"]
        marking_id = self.tlp_marking

        # Vulnerability
        stix_vuln = stix2.Vulnerability(
            id=Vulnerability.generate_id(vuln_cve),
            name=f"{vuln_cve}",
            description=f"{description}",
            created_by_ref=self.created_by_stix["id"],
            created=f"{created}",
            object_marking_refs=[f"{marking_id}"],
        )
        stix_objects.append(stix_vuln)
        vuln_id = stix_vuln["id"]

        # Software vendor
        stix_org = stix2.Identity(
            id=Identity.generate_id(vendor_name, "organization"),
            name=f"{vendor_name}",
            identity_class="organization",
            description="Software Vendor",
            created=f"{created}",
            created_by_ref=f"{created_by_id}",
            object_marking_refs=[f"{marking_id}"],
            allow_custom=True,
            custom_properties={"x_opencti_organization_type": "vendor"},
        )
        stix_objects.append(stix_org)
        org_id = stix_org["id"]

        stix_software = stix2.Software(
            name=f"{product}",
            vendor=f"{vendor_name}",
            allow_custom=True,
            custom_properties={
                "created_by_ref": created_by_id,
            },
        )
        stix_objects.append(stix_software)
        software_id = stix_software["id"]

        software_vendor_relationship = stix2.Relationship(
            id=StixCoreRelationship.generate_id(
                "related-to", software_id, org_id, created
            ),
            relationship_type="related-to",
            description="maintains",
            start_time=f"{created}",
            source_ref=f"{software_id}",
            target_ref=f"{org_id}",
            confidence="100",
            created_by_ref=f"{created_by_id}",
            object_marking_refs=[f"{marking_id}"],
        )
        stix_objects.append(software_vendor_relationship)

        # Infrastructures
        if self.cisa_create_infrastructures:
            stix_infrastructure = stix2.Infrastructure(
                id=Infrastructure.generate_id(product),
                name=f"{product}",
                created_by_ref=self.created_by_stix["id"],
                created=f"{created}",
                object_marking_refs=[f"{marking_id}"],
            )
            stix_objects.append(stix_infrastructure)
            infra_id = stix_infrastructure["id"]
            infra_vuln_relationship = stix2.Relationship(
                id=StixCoreRelationship.generate_id("has", infra_id, vuln_id, created),
                relationship_type="has",
                source_ref=f"{infra_id}",
                target_ref=f"{vuln_id}",
                start_time=f"{created}",
                confidence="100",
                object_marking_refs=[f"{marking_id}"],
            )
            stix_objects.append(infra_vuln_relationship)
            infra_vendor_relationship = stix2.Relationship(
                id=StixCoreRelationship.generate_id(
                    "related-to", org_id, infra_id, created
                ),
                relationship_type="related-to",
                description="maintains",
                source_ref=f"{org_id}",
                target_ref=f"{infra_id}",
                start_time=f"{created}",
                confidence="100",
                object_marking_refs=[f"{marking_id}"],
            )
            stix_objects.append(infra_vendor_relationship)

        software_vuln_relationship = stix2.Relationship(
            id=StixCoreRelationship.generate_id("has", software_id, vuln_id, created),
            relationship_type="has",
            source_ref=f"{software_id}",
            target_ref=f"{vuln_id}",
            start_time=f"{created}",
            confidence="100",
            created_by_ref=f"{created_by_id}",
            object_marking_refs=[f"{marking_id}"],
        )
        stix_objects.append(software_vuln_relationship)

        bundle = stix2.Bundle(objects=stix_objects, allow_custom=True).serialize()
        self.helper.log_info("CISA Bundle Complete")
        return bundle

    def process_data(self):
        try:
            # Get the current timestamp and check
            timestamp = int(time.time())
            current_state = self.helper.get_state()
            if current_state is not None and "last_run" in current_state:
                last_run = current_state["last_run"]
                self.helper.log_info(
                    "Connector last run: "
                    + datetime.datetime.utcfromtimestamp(last_run).strftime(
                        "%Y-%m-%d %H:%M:%S"
                    )
                )
            else:
                last_run = None
                self.helper.log_info("Connector has never run")
            # If the last_run is more than interval-1 day
            if last_run is None or (
                (timestamp - last_run) > ((int(self.cisa_interval) - 1) * 60 * 60 * 24)
            ):
                self.helper.log_info("Connector will run!")

                now = datetime.datetime.utcfromtimestamp(timestamp)
                friendly_name = "CISA run @ " + now.strftime("%Y-%m-%d %H:%M:%S")
                work_id = self.helper.api.work.initiate_work(
                    self.helper.connect_id, friendly_name
                )
                if self.cisa_catalog_url is not None and len(self.cisa_catalog_url) > 0:
                    cisa_data = self.retrieve_data(self.cisa_catalog_url)
                    self.set_created_by_stix(org=self.org)
                    self.set_tlp_marking(tlp_mark=self.tlp)
                    cisa_data = json.loads(cisa_data)
                    for vuln in cisa_data["vulnerabilities"]:
                        bundle = self.build_bundle(vuln)
                        self.send_bundle(work_id, bundle)

                # Store the current timestamp as a last run
                message = "Connector successfully run, storing last_run as " + str(
                    timestamp
                )
                self.helper.log_info(message)
                self.helper.set_state({"last_run": timestamp})
                self.helper.api.work.to_processed(work_id, message)
                self.helper.log_info(
                    "Last_run stored, next run in: "
                    + str(round(self.get_interval() / 60 / 60 / 24, 2))
                    + " days"
                )
            else:
                new_interval = self.get_interval() - (timestamp - last_run)
                self.helper.log_info(
                    "Connector will not run, next run in: "
                    + str(round(new_interval / 60 / 60 / 24, 2))
                    + " days"
                )
        except (KeyboardInterrupt, SystemExit):
            self.helper.log_info("Connector stop")
            sys.exit(0)
        except Exception as e:
            self.helper.log_error(str(e))

    def run(self):
        self.helper.log_info("Fetching CISA Known Exploited Vulnerabilities...")
        get_run_and_terminate = getattr(self.helper, "get_run_and_terminate", None)
        if callable(get_run_and_terminate) and self.helper.get_run_and_terminate():
            self.process_data()
            self.helper.force_ping()
        else:
            while True:
                self.process_data()
                time.sleep(60)

    def send_bundle(self, work_id: str, serialized_bundle: str) -> None:
        try:
            self.helper.send_stix2_bundle(
                serialized_bundle,
                update=self.update_existing_data,
                work_id=work_id,
            )
        except Exception as e:
            self.helper.log_error(f"Error while sending bundle: {e}")


if __name__ == "__main__":
    try:
        connector = Cisa()
        connector.run()
    except Exception as e:
        print(e)
        time.sleep(10)
        exit(0)
